/*
 * PennTreeReader.java
 *
 * Created on July 12, 2006, 11:40 AM
 *
 */

package galronnlp.io;

import galronnlp.perceptron.FeatureVector;
import galronnlp.perceptron.HeadFinderRules;
import galronnlp.perceptron.Node;
import galronnlp.perceptron.Grammar;
import galronnlp.util.BinaryTree;
import galronnlp.util.LexicalTree;
import galronnlp.util.NaryTree;
import galronnlp.util.ProductionGenerator;
import galronnlp.util.Symbol;
import galronnlp.util.SymbolGenerator;
import galronnlp.util.Tree;
import galronnlp.util.Tuple;
import galronnlp.util.UnaryTree;
import java.io.IOException;
import java.io.Reader;
import java.io.BufferedReader;
import java.io.StringReader;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Scanner;
import java.util.regex.*;
/**
 *
 * @author Daniel A. Galron
 */
public class PennTreeReader {
    
    private BufferedReader reader;
    private SymbolGenerator gen;
    private FeatureVector parameters = null;
    private boolean next = true;
    
    /**
     * Creates a new instance of PennTreeReader
     */
    public PennTreeReader(BufferedReader r, SymbolGenerator g, FeatureVector parameters) {
        reader = r;
        gen = g;
        parameters = parameters;
    }

    public PennTreeReader(BufferedReader r, SymbolGenerator g) {
        reader = r;
        gen = g;
    }    
    
    public PennTreeReader(BufferedReader r) {
        reader = r;
    }
    
    public boolean hasNext() {
        return next;
    }
    
    public void close() {
        try {
            reader.close();
        } catch (IOException ex) {
            ex.printStackTrace();
        }
    }
    
    /**
     * Read in the next tree from the file. These trees must be in the Penn Treebank
     * mrg format.
     */
    public Tree readTree() {
        // Get the tree from teh treebank represented as a string
        String treeString = this.readNextTreeString();
        Tree ret = null;
        // If the string is not empty (that is, if the string is not an empty line
        // read from the file, construct the tree and return it.
        if(!treeString.equals("")) {
            try {
                ret = this.constructTree(treeString);
            } catch (malformedTreeException ex) {
                ex.printStackTrace();
            }
        }
        return ret;
    }
    
    public Node readNodeTree() {
        String treeString = this.readNextTreeString();
        Node ret = null;
        if(!treeString.equals("")) {
            try {
                LinkedList<Tuple> inputSentence = this.getTaggedSentence(new String(treeString));
                ret = this.constructNodeTree(treeString, null, inputSentence);
            } catch(malformedTreeException ex) {
                ex.printStackTrace();
            }
        }
        return ret;
    }
    
    public Node addStartSymbol(Node tree, Symbol startSymbol, Symbol endSymbol, Symbol endSymbolPOS) {
        Node n = new Node(null, Node.NONTERMINAL, startSymbol, endSymbol, endSymbolPOS, gen, null);
        //System.err.println(n);
        //System.err.println(tree);
        n.addChild(tree);
        Node lastChain = new Node(null, Node.POS, endSymbolPOS, null, null, gen, null);
        lastChain.addChild(new Node(null, Node.TERMINAL, endSymbol, null, null, gen, null));
        n.addChild(lastChain);
        return n;
    }
    
    private Node constructNodeTree(String treeString, Node parent, LinkedList<Tuple> inputSentence) throws malformedTreeException {
        // Trim whitespace from the tree String
        treeString = treeString.trim();
        // Create a new tree scanner to build this part of the tree
        PennTreeScanner ts = new PennTreeScanner(treeString, gen);
        // If the next part of the tree is a leaf (that is, it is of format:
        // (X y), then:
        if(ts.hasNextLeaf()) {
            // Trim the parentheses from this leaf
            String noParenTreeString = this.trimParens(ts.nextLeaf());
            // Create a tree scanner specially for the String "X y"
            // INVARIANT: noParenTreeString is of format "X y"
            PennTreeScanner ls = new PennTreeScanner(noParenTreeString, gen);
            // get the POS and the lexical item
            Symbol LHS = ls.nextSymbol();
            Symbol RHS = ls.nextSymbol();
            Tuple lookAhead = inputSentence.removeFirst();
            Node nonterminal = new Node(parent, Node.POS, LHS, (Symbol)lookAhead.get(0), (Symbol)lookAhead.get(1), gen, parameters);
            Node terminal = new Node(nonterminal, Node.TERMINAL, RHS, (Symbol)lookAhead.get(0), (Symbol)lookAhead.get(1), gen, parameters);
            nonterminal.addChild(terminal);
            return nonterminal;
        // Otherwise, if the next part of the tree is of type constituent, that is,
        // if it is a full subtree, containing one or more other subtrees or lexical
        // items, (this corresponds to a production in the recursively built tree)
        } else if(ts.hasNextConstituent()) {
            // Trim the parentheses from this subtree string
            String noParenTreeString = this.trimParens(ts.nextConstituent());
            PennTreeScanner cs = new PennTreeScanner(noParenTreeString, gen);
            //INVARIANT: noParenTreeString = LHS (.*)
            Symbol LHS = cs.nextSymbol();
            Tuple lookAhead = inputSentence.getFirst();
            Node nonterminal = new Node(parent, Node.NONTERMINAL, LHS, (Symbol)lookAhead.get(0), (Symbol)lookAhead.get(1), gen, parameters);
            while(cs.hasNextConstituent()) {
                String next = cs.nextConstituent();
                Node child = this.constructNodeTree(next, nonterminal, inputSentence);
                nonterminal.addChild(child);
            }
            return nonterminal;
        // Deals with the case where the next thing read is a pair of open parentheses,
        // signifying (in the case of the WSJ corpus) that the treeString is the top
        // level of the tree.
        } else if(ts.hasNextStart()) {
            String noParenTreeString = this.trimParens(ts.nextStart());
            Symbol StartSymbol = gen.create("S!");
            Symbol EndSymbolPOS = gen.create("STOP");
            Symbol EndSymbol = gen.create("/S");
            Tuple lookAhead = inputSentence.getFirst();
            Node startNode = new Node(parent, Node.NONTERMINAL, StartSymbol, (Symbol)lookAhead.get(0), (Symbol)lookAhead.get(1), gen, parameters);
            Node lc = this.constructNodeTree(noParenTreeString.trim(), startNode, inputSentence);
            Node endNodePOS = new Node(startNode, Node.POS, EndSymbolPOS, null, null, gen, parameters);
            Node endNode = new Node(endNodePOS, Node.TERMINAL, EndSymbol, null, null, gen, parameters);
            endNodePOS.addChild(endNode);
            startNode.addChild(lc);
            startNode.addChild(endNodePOS);
            return startNode;
        } else {
            throw new malformedTreeException(treeString);
        }        
    }

    private LinkedList<Tuple> getTaggedSentence(String treeString) throws malformedTreeException {
        treeString = treeString.trim();
        PennTreeScanner ts = new PennTreeScanner(treeString, gen);
        LinkedList<Tuple> ret = new LinkedList<Tuple>();
        if(ts.hasNextLeaf()) {
            String noParenTreeString = this.trimParens(ts.nextLeaf());
            PennTreeScanner ls = new PennTreeScanner(noParenTreeString, gen);
            Symbol POS  = ls.nextSymbol();
            Symbol word = ls.nextSymbol();
            Tuple r = new Tuple(2);
            r.put(0, word);
            r.put(1, POS);
            ret.addLast(r);
            return ret;
        } else if(ts.hasNextConstituent()) {
            String noParenTreeString = this.trimParens(ts.nextConstituent());
            PennTreeScanner cs = new PennTreeScanner(noParenTreeString, gen);
            Symbol throwaway = cs.nextSymbol();
            while(cs.hasNextConstituent()) {
                String next = cs.nextConstituent();
                ret.addAll(this.getTaggedSentence(next));
            }
            return ret;
        } else if(ts.hasNextStart()) {
            String noParenTreeString = this.trimParens(ts.nextStart());
            return this.getTaggedSentence(noParenTreeString.trim());            
        } else {
            throw new malformedTreeException(treeString);
        }
    }
    
    private Tree constructTree(String treeString) throws malformedTreeException {
        // Trim whitespace from the tree String
        treeString = treeString.trim();
        // Create a new tree scanner to build this part of the tree
        PennTreeScanner ts = new PennTreeScanner(treeString, gen);
        // If the next part of the tree is a leaf (that is, it is of format:
        // (X y), then:
        if(ts.hasNextLeaf()) {
            // Trim the parentheses from this leaf
            String noParenTreeString = this.trimParens(ts.nextLeaf());
            // Create a tree scanner especially for the String "X y"
            // INVARIANT: noParenTreeString is of format "X y"
            PennTreeScanner ls = new PennTreeScanner(noParenTreeString, gen);
            // get the category and lexical item
            Symbol LHS = ls.nextSymbol();
            Symbol RHS = ls.nextSymbol();
            // Return a new LexicalTree item containing the LHS and RHS
            return new LexicalTree(LHS, RHS);
        // Otherwise, if the next part of tree is of type constituent, that is,
        // if it is a full subtree, containing one or more other subtrees or lexical items
        // (this corresponds to a production in the recursively built tree.
        } else if(ts.hasNextConstituent()) {
            // Trim the parentheses from this subtree string
            String noParenTreeString = this.trimParens(ts.nextConstituent());
            // Create a tree scanner for this subtree
            PennTreeScanner cs = new PennTreeScanner(noParenTreeString, gen);
            // INVARIANT: noParenTreeString = LHS (.*)
            Symbol LHS = cs.nextSymbol();
            LinkedList<Tree> children = new LinkedList<Tree>();
            // for each subsequent subtree which is a child of this level in the tree
            while(cs.hasNextConstituent()) {
                String next = cs.nextConstituent();
                // Recursively construct the tree representation for that subtree, 
                // and add it to the list of children of this side.
                children.addLast(this.constructTree(next.trim()));
            }
            if(children.size() == 1) {
                return new UnaryTree(LHS, children.get(0));
            }
            else if(children.size() == 2)
                return new BinaryTree(LHS, children.get(0), children.get(1));
            else
                return new NaryTree(LHS, children);
        // Deals with the case where the next thing read is a pair of open parentheses,
        // signifying (in the case of the WSJ corpus) that the treeString is the top
        // level of the tree.
        } else if(ts.hasNextStart()) {
            String noParenTreeString = this.trimParens(ts.nextStart());
            return this.constructTree(noParenTreeString.trim());
        } else {
            throw new malformedTreeException(treeString);
        }
    }
    
    private String trimParens(String string) {
        int first = string.indexOf('(') + 1;
        int last = string.lastIndexOf(')');
        return string.substring(first, last).trim();
    }
    
    public static void main(String[] args) {
        String tree = "( (S\n" +
                      "   (NP-SBJ\n" +
                      "     (NP (NNP Pierre) (NNP Vinken) )\n" +
                      "     (, ,)\n" +
                      "     (ADJP\n" +
                      "       (NP (CD 61) (NNS years) )\n" +
                      "       (JJ old) )\n" +
                      "     (, ,) )\n" +
                      "   (VP (MD will)\n" +
                      "     (VP (VB join)\n" +
                      "       (NP (DT the) (NN board) )\n" +
                      "       (PP-CLR (IN as)\n" +
                      "         (NP (DT a) (JJ nonexecutive) (NN director) ))\n" +
                      "       (NP-TMP (NNP Nov.) (CD 29) )))\n" +
                      "     (. .) ))";
        //String tree2 = "(NP dog)";
        SymbolGenerator gen = new SymbolGenerator();
        PennTreeReader r = new PennTreeReader(new BufferedReader(new StringReader(tree)), gen, null);
        HeadFinderRules rules = new HeadFinderRules(gen);
        Grammar g = new Grammar(new ProductionGenerator(), gen);
        Node n = r.readNodeTree();
        //System.out.println(n);
        n.fireFeatures(rules);
        n.addToGrammar(g);
        
        String XML = "<?xml version=\"1.0\"?>\n" +
                     "<grammar type=\"perceptron\">\n" +
                     n.subtreeFeatures.getXMLRep() + 
                     "<productions>\n" +
                     n.extractXMLGrammarRep() +
                     "</productions>\n" +
                     "</grammar>\n";
        System.out.println(g.getXMLRep());
        //System.out.println(g);
        //Node copy = n.copy();
        //System.out.println(copy);
        //System.out.println(XML);
    }
    
   
    private String readNextTreeString() {
        String tree = "";
        int parensSeen = 0;
        try {
            // Read the first line in the reader, and count the number of open parentheses.
            String line = reader.readLine();
            // Skip over blank lines
            while(line.length() == 0) {
                line = reader.readLine();
            }
            for(int i = 0; i < line.length(); i++) {
                if(line.charAt(i) == '(') parensSeen++;
                else if(line.charAt(i) == ')') parensSeen--;
            }
            tree = line.trim();
        } catch (IOException ex) {
            next = false;
            return "";
        } catch(NullPointerException ex) {
            // Exception is thrown when the reader has reached the end of file.
            // This means that there aren't any trees left in the file, which causes
            // hasNext() to be false
            next = false;
            return "";
        }        
        // Read in the next lines until close parentheses have been found for all open
        // parentheses
        while(parensSeen != 0) {
            try {
                String line = reader.readLine();
                for(int i = 0; i < line.length(); i++) {
                    if(line.charAt(i) == '(') parensSeen++;
                    else if(line.charAt(i) == ')') parensSeen--;
                }
                tree += line;
            } catch(IOException ex) {
                next = false;
                return "";
            } catch(NullPointerException ex) {
                next = false;
                return "";
            }
        }
        return tree + "\n";
    }
    
}
